import torch
import torch.nn as nn
import torchvision
from torchvision.models.feature_extraction import create_feature_extractor

__all__ = ["resnet18"]

def resnet18(pretrained,**kwargs):
    weight = None
    if pretrained is not None:
        weight = torchvision.models.ResNet18_Weights
    model = torchvision.models.resnet18(weights=weight)
    node_name = "layer4.1.relu_1"
    model = create_feature_extractor(model,{node_name:'output'})
    return model

def build_backbone(model_name='resnet18',pretrained=True):
    if model_name =='resnet18':
        fea_dims = 512
        return resnet18(pretrained),fea_dims
    else:
        raise NotImplementedError("backbone {} not implemented".format(model_name))

